import torch
import torch.nn as nn
import numpy as np


class ANFISLayer(nn.Module):
    def __init__(self, input_dim, S=2, T=16, learning_rate=0.01):
        super(ANFISLayer, self).__init__()
        self.S = S
        self.T = T
        self.learning_rate = learning_rate

        # Initialize parameters
        self.c1 = nn.Parameter(torch.empty(S, input_dim).uniform_(-1, 1))
        self.b1 = nn.Parameter(torch.empty(S, input_dim).uniform_(-1, 1))
        self.w1 = nn.Parameter(torch.empty(1, T).uniform_(0, 1))

    def forward(self, samplein):
        # Normalization
        samplein_min = samplein.min(dim=1, keepdim=True).values
        samplein_max = samplein.max(dim=1, keepdim=True).values
        sampleinnorm = (samplein - samplein_min) / (samplein_max - samplein_min)

        # Initialize output tensor
        Y = torch.zeros_like(sampleinnorm)

        # Compute membership degrees
        for m in range(sampleinnorm.shape[0]):
            u1 = torch.exp(-((sampleinnorm[m].unsqueeze(0) - self.c1) ** 2) / (self.b1 ** 2))

            # Rule generation layer
            alpha1 = torch.zeros(self.T, device=sampleinnorm.device)
            for i in range(self.S):
                a = 1 if i != 1 else 8
                for p in range(self.S):
                    b = 1 if p != 1 else 4
                    for q in range(self.S):
                        c = 1 if q != 1 else 2
                        for k in range(self.S):
                            idx = i * a + p * b + q * c + k
                            if idx < self.T:  # Ensure idx is within bounds
                                alpha1[idx] = u1[i, 0] * u1[p, 0] * u1[q, 0] * u1[
                                    k, 0]  # Index to get the first element

            alphasum = alpha1.sum()
            Y[m] = torch.dot(self.w1.squeeze(), alpha1) / alphasum if alphasum != 0 else 0

        return Y


# Example usage

import torch

# Set random seed for reproducibility
torch.manual_seed(0)

# Generate random input data
n_samples = 24  # Number of samples
m_features = 64   # Number of features

# Sample input (samplein) and output (sampleout)
samplein = torch.rand(n_samples, m_features)  # Random values between 0 and 10
sampleout = torch.rand(n_samples, m_features)  # Random values between 0 and 10

# Create ANFIS layer
input_dim = samplein.shape[1]  # Input feature dimension
print(samplein.shape)
print(input_dim)

anfis_layer = ANFISLayer(input_dim)

# Forward pass
output = anfis_layer(samplein)

# Print output
print("Output:", output)